{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "import torch\n",
    "import string\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "\n",
    "def is_chinese(text):\n",
    "    \"\"\"判断文本是否含中文字符\"\"\"\n",
    "    return any('\\u4e00' <= ch <= '\\u9fff' for ch in text)\n",
    "\n",
    "def normalize(s):\n",
    "    def remove_articles(text):\n",
    "        return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\n",
    "\n",
    "    def white_space_fix(text):\n",
    "        return \" \".join(text.split())\n",
    "\n",
    "    def remove_punc(text):\n",
    "        chinese_punc = \"！？｡＂＃＄％＆＇（）＊＋，－．／：；＜＝＞＠［＼］＾＿｀｛｜｝～“”‘’、。：《》【】\"\n",
    "        exclude = set(string.punctuation + chinese_punc)\n",
    "        return \"\".join(ch for ch in text if ch not in exclude)\n",
    "\n",
    "    def lower(text):\n",
    "        return text.lower()\n",
    "\n",
    "    s = lower(s)\n",
    "    s = remove_punc(s)\n",
    "\n",
    "    if is_chinese(s):\n",
    "        s = s.replace(\" \", \"\")  # 中文一般去除所有空白\n",
    "    else:\n",
    "        s = remove_articles(s)\n",
    "        s = white_space_fix(s)\n",
    "\n",
    "    return s\n",
    "\n",
    "def compute_f1(prediction, ground_truth):\n",
    "    if prediction is None:\n",
    "        return 0.0\n",
    "\n",
    "    norm_pred = normalize(prediction)\n",
    "    norm_gt = normalize(ground_truth)\n",
    "\n",
    "    # 中文使用字符级，英文使用词级\n",
    "    if is_chinese(norm_pred) or is_chinese(norm_gt):\n",
    "        pred_tokens = list(norm_pred)\n",
    "        gt_tokens = list(norm_gt)\n",
    "    else:\n",
    "        pred_tokens = norm_pred.split()\n",
    "        gt_tokens = norm_gt.split()\n",
    "\n",
    "    common = set(pred_tokens) & set(gt_tokens)\n",
    "    num_same = len(common)\n",
    "\n",
    "    if num_same == 0:\n",
    "        return 0.0\n",
    "\n",
    "    precision = num_same / len(pred_tokens)\n",
    "    recall = num_same / len(gt_tokens)\n",
    "    f1 = 2 * precision * recall / (precision + recall)\n",
    "    return f1\n",
    "\n",
    "def exact_match_score(prediction, ground_truth):\n",
    "    if prediction is None:\n",
    "        return 0.0\n",
    "    return int(normalize(prediction) == normalize(ground_truth))\n",
    "\n",
    "def plot_images(image_paths):\n",
    "    num_images = len(image_paths)\n",
    "    \n",
    "    fig, axes = plt.subplots(1, num_images, figsize=(5 * num_images, 5))\n",
    "    \n",
    "    for i, image_path in enumerate(image_paths):\n",
    "        img = mpimg.imread(image_path)\n",
    "\n",
    "        # 如果是灰度图（2D），扩展为 RGB\n",
    "        if img.ndim == 2:\n",
    "            img = np.stack([img]*3, axis=-1)\n",
    "        # 如果是带 alpha 通道的 RGBA 图，去掉 alpha\n",
    "        elif img.shape[2] == 4:\n",
    "            img = img[:, :, :3]\n",
    "\n",
    "        ax = axes if num_images == 1 else axes[i]\n",
    "        ax.imshow(img)\n",
    "        ax.set_title(f'Image {i+1}')\n",
    "        ax.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def extract_coordinates(text):\n",
    "    pattern = r'\\[\\s*(\\d+)\\s*,\\s*(\\d+)\\s*,\\s*(\\d+)\\s*,\\s*(\\d+)\\s*\\]'\n",
    "    matches = re.findall(pattern, text)\n",
    "    # 转换为整数列表格式\n",
    "    coordinates = [list(map(int, match)) for match in matches]\n",
    "    return coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_path = \"./evaluation_results/3b_mat_coding_grpo_step1400.json\"\n",
    "with open(json_path, \"r\") as f:\n",
    "    data = json.load(f)\n",
    "print(len(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = 0 \n",
    "em = 0\n",
    "f1_all = []\n",
    "em_all = []\n",
    "count = 0\n",
    "none_count = 0\n",
    "for item in data[:]:\n",
    "    if 'pred_answer' != None:\n",
    "        count += 1\n",
    "        # 计算 f1 和 em \n",
    "        pred = item['pred_answer']\n",
    "        gts = item['gt']\n",
    "        # 若gt是str，统一转换为列表处理\n",
    "        if isinstance(gts, str):\n",
    "            gts = [gts]\n",
    "        f1 = max([compute_f1(pred, gt) for gt in gts])\n",
    "        em = max([exact_match_score(pred, gt) for gt in gts])\n",
    "        if em == 1:\n",
    "            f1 =1\n",
    "        # print(\"F1: \" + str(f1))\n",
    "        # print(\"EM: \" + str(em))\n",
    "        f1_all.append(f1)\n",
    "        em_all.append(em)\n",
    "    else:\n",
    "        count += 1\n",
    "        none_count += 1\n",
    "        f1 = 0.0\n",
    "        em = 0.0\n",
    "        # print(\"F1: \" + str(f1))\n",
    "        # print(\"EM: \" + str(em))\n",
    "        f1_all.append(f1)\n",
    "        em_all.append(em)\n",
    "\n",
    "\n",
    "print(count)\n",
    "print(none_count)\n",
    "print('Simple F1:', sum(f1_all[:70])/70)\n",
    "print('Simple EM:', sum(em_all[:70])/70)\n",
    "print('Hard F1:', sum(f1_all[70:])/130)\n",
    "print('Hard EM:', sum(em_all[70:])/130)\n",
    "print('All F1:', sum(f1_all)/200)\n",
    "print('All EM:', sum(em_all)/200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
